# In ML, we always use yml file to define the structure of Model.
# Please define a function called CreateModel for me to create different model.
# At the last the function, using json to store the dictionary to a "config.json"
# My yaml file is define by:
#   1. first line: type of model (assume we only have two types: CV model, NLP model)
#   2. The rest would be model config
# You already have two function: CreateCVModel/ CreateaNLPModel, you can directly use them
# CreateCVModel/ CreateaNLPModel is well-defined and doesn't have any other variable input. So please not input "type" into that.
# You can also use pyyaml for help


import json
import yaml


# 这个CreateModel函数首先读取YAML文件,根据type字段判断是创建CV还是NLP模型,然后调用相应的函数创建模型对象。最后将模型配置以字典形式存储到config.json文件中。
# 读取到config.json的model config内容不应该包含 type 属性内容，而且需要美化过的json 格式，不需要一行打印
# 因此在构造model_config时,我用model_def.copy()创建一个副本,这样可以删除type字段而不影响原字典。
# 并使用json库的indent参数来格式化JSON输出,使其格式是美化过的
# CreateCVModel和CreateNLPModel函数内部我先只是执行打印,以题目前提为主，假设其中内部实现已完成
def CreateModel(yaml_file):
    with open(yaml_file) as f:
        model_def = yaml.safe_load(f)

    if model_def['type'] == 'CV':
        print(model_def)
        CreateCVModel()
    elif model_def['type'] == 'NLP':
        print(model_def)
        CreateNLPModel()

    model_config = model_def.copy()
    del model_config['type']

    json_content = {
        'type': model_def['type'],
        'config': model_config
    }

    with open('config.json', 'w') as f:
        json.dump(json_content, f, indent=4)


def CreateCVModel():
    # 创建CV模型,返回模型对象
    print("CreateCVModel")


def CreateNLPModel():
    # 创建NLP模型,返回模型对象
    print("CreateNLPModel")


if __name__ == "__main__":
    CreateModel("CV_yaml.yaml")
